#include "JointTrackingController.h"
#include "Modeling/DynamicPath.h"


JointTrackingController::JointTrackingController(Robot& _robot)
  :RobotController(_robot)
{
  qdesDefault = _robot.q;
}
    
//subclasses should override this
void JointTrackingController::GetDesiredState(Config& q_des,Vector& dq_des)
{
  q_des = qdesDefault;
  dq_des.setZero();
}

void JointTrackingController::Update(Real dt)
{
  Assert(command != NULL);

  Config qdes(robot.links.size()),dqdes(robot.links.size());
  GetDesiredState(qdes,dqdes);
  for(size_t i=0;i<robot.drivers.size();i++) {
    if(robot.drivers[i].type == RobotJointDriver::Normal) {
      command->actuators[i].SetPID(qdes(robot.drivers[i].linkIndices[0]),dqdes(robot.drivers[i].linkIndices[0]),command->actuators[i].iterm);
    }
    else {
      robot.q = qdes;
      robot.dq = dqdes;
      //printf("Desired affine driver value %g, vel %g\n",robot.GetDriverValue(i),robot.GetDriverVelocity(i));
      command->actuators[i].SetPID(robot.GetDriverValue(i),robot.GetDriverVelocity(i),command->actuators[i].iterm);
    }
  }
  RobotController::Update(dt);
}

void JointTrackingController::Reset()
{
  if(command) {
    for(size_t i=0;i<command->actuators.size();i++) 
      command->actuators[i].SetPID(0,0);
    command->ResetPIDIntegrals();
  }
  RobotController::Reset();
}

MilestoneTrackingController::MilestoneTrackingController(Robot& _robot)
  :JointTrackingController(_robot),pathParameter(0),velScale(1.0),accScale(1.0),modifySpeedByError(false),modifySpeedCoeff(50.0)
{
  xcur=_robot.q;
  dxcur=_robot.dq;
  ramp.endTime=0;
  SetMilestone(_robot.q);
}


void MilestoneTrackingController::SetMilestone(const Vector& x,const Vector& dx)
{
  Real oldPathParameter = pathParameter;
  ParabolicRampND oldRamp = ramp;
  pathParameter=0;
  ramp.x0=xcur;
  ramp.dx0=dxcur;
  ramp.x1=x;
  ramp.dx1=dx;
  Vector velMax,accMax;
  velMax.mul(robot.velMax,velScale);
  accMax.mul(robot.accMax,accScale);
  bool res=ramp.SolveMinTime(accMax,velMax);
  if(!res) {
    pathParameter = oldPathParameter;
    ramp = oldRamp;
  }
}

void MilestoneTrackingController::SetMilestone(const Vector& x)
{
  ParabolicRampND oldRamp = ramp;
  Real oldPathParameter = pathParameter;
  pathParameter=0;
  ramp.x0=xcur;
  ramp.dx0=dxcur;
  ramp.x1=x;
  ramp.dx1.resize(x.n); ramp.dx1.setZero();
  Vector velMax,accMax;
  velMax.mul(robot.velMax,velScale);
  accMax.mul(robot.accMax,accScale);
  bool res=ramp.SolveMinTime(accMax,velMax);
  if(!res) {
    ramp = oldRamp;
    pathParameter = oldPathParameter;
  }
}

void MilestoneTrackingController::SetMilestone(const Vector& x,Real minTime)
{
  pathParameter=0;
  ramp.x0=xcur;
  ramp.dx0=dxcur;
  ramp.x1=x;
  ramp.dx1.resize(x.n); ramp.dx1.setZero();
  Vector velMax;
  velMax.mul(robot.velMax,velScale);
  bool res=ramp.SolveMinAccel(velMax,minTime);
  if(!res) {
    Vector accMax;
    accMax.mul(robot.accMax,accScale);
    res=ramp.SolveMinTime(accMax,velMax);
    Assert(res);
  }
}

void MilestoneTrackingController::SetLinearVelocity(const Vector& dx,Real t)
{
  pathParameter = 0;
  ramp.SetLinear(xcur,xcur+dx*t,t);
}

Real MilestoneTrackingController::GetSpeedScale(const Config& q_des,const Vector& dq_des) const
{
  if(modifySpeedByError) {
    //adjust the path speed by the tracking error
    Real trackingError=0;
    for(size_t i=0;i<robot.joints.size();i++) {
      if(robot.joints[i].type == RobotJoint::Normal) {
	int k=robot.joints[i].linkIndex;
	trackingError += Sqr(q_des(k) - sensors->q(k));
      }
    }
    return One/(One+modifySpeedCoeff*trackingError);
  }
  return 1.0;
}
    
void MilestoneTrackingController::GetDesiredState(Config& q_des,Vector& dq_des)
{
  Config qNext;
  ramp.Evaluate(pathParameter,q_des);
  ramp.Derivative(pathParameter,dq_des);

  if(pathParameter >= ramp.endTime) {
    q_des = ramp.x1;
    dq_des = ramp.dx1;
  }
  else {
    Real speedModifier = GetSpeedScale(xcur,dxcur);
    dq_des *= speedModifier;  
  }
  xcur=q_des;
  dxcur=dq_des;
}

void MilestoneTrackingController::Update(Real dt)
{
  Real speedModifier = GetSpeedScale(xcur,dxcur);
  pathParameter += dt * speedModifier;
  if(pathParameter >= ramp.endTime) {
    pathParameter = ramp.endTime;
  }
  JointTrackingController::Update(dt);
}

void MilestoneTrackingController::Reset()
{
  pathParameter = 0;
  //xcur = xcur;
  dxcur.setZero();
  SetMilestone(xcur);
  JointTrackingController::Reset();
}

bool MilestoneTrackingController::ReadState(File& f)
{
  if(!JointTrackingController::ReadState(f)) return false;
  if(!ReadFile(f,xcur)) return false;
  if(!ReadFile(f,dxcur)) return false;
  if(!ReadFile(f,ramp.x0)) return false;
  if(!ReadFile(f,ramp.x1)) return false;
  if(!ReadFile(f,ramp.dx0)) return false;
  if(!ReadFile(f,ramp.dx1)) return false;
  if(!ReadFile(f,pathParameter)) return false;
  if(!ReadFile(f,velScale)) return false;
  if(!ReadFile(f,accScale)) return false;
  if(!ReadFile(f,modifySpeedByError)) return false;
  if(!ReadFile(f,modifySpeedCoeff)) return false;
  //setup ramp
  Vector velMax,accMax;
  velMax.mul(robot.velMax,velScale);
  accMax.mul(robot.accMax,accScale);
  bool res=ramp.SolveMinTime(accMax,velMax);
  Assert(res);
  return true;
}

bool MilestoneTrackingController::WriteState(File& f) const
{
  if(!JointTrackingController::WriteState(f)) return false;
  if(!WriteFile(f,xcur)) return false;
  if(!WriteFile(f,dxcur)) return false;
  if(!WriteFile(f,ramp.x0)) return false;
  if(!WriteFile(f,ramp.x1)) return false;
  if(!WriteFile(f,ramp.dx0)) return false;
  if(!WriteFile(f,ramp.dx1)) return false;
  if(!WriteFile(f,pathParameter)) return false;
  if(!WriteFile(f,velScale)) return false;
  if(!WriteFile(f,accScale)) return false;
  if(!WriteFile(f,modifySpeedByError)) return false;
  if(!WriteFile(f,modifySpeedCoeff)) return false;
  return true;
}


